# This script relies on some external programs. The required Ubuntu
# packages are dvgrab, tovid, imagemagick (for convert), smilutils
# (for image2raw) and transcode (for tcprobe).

from __future__ import with_statement
import sys, os
import subprocess
import time
import re
import pickle

_numFramesForDate = 50

class Error(Exception): pass
    
def log(message):
    print >> sys.stderr, message
    
def call(command):
    log('Executing: %s' % command)
    return subprocess.check_call(command, shell=True)

def driveSpaceInGb(path):
    s = os.statvfs(path)
    return s.f_bsize * s.f_bfree / 1024**3

class DvFile(object):
    def __init__(self, filename):
        self.filename = filename
        self.isValid = True
        self.needsTimestamp = False
        # Check the validity of the filename.
        try:
            self.TimeTuple()
        except ValueError:
            raise Error('Filename does not match the dvgrab filename pattern.', filename)

    def __repr__(self):
        return 'DvFile(%s)' % repr(self.__dict__)

    def ReportString(self):
        return '%s IsValid=%5s NeedsTimestamp=%5s NumFrames=%5d' % (self.filename, self.isValid, self.needsTimestamp, self.NumFrames())
    
    def NumFrames(self):
        process = subprocess.Popen(('tcprobe -i %s' % self.filename).split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout = process.communicate()[0]
        match = re.search('length: ([0-9]*) frames', stdout)
        if match is None:
            raise Error('Unexpected output from tcprobe.', self.filename)
        return int(match.group(1))
    
    def Time(self):
        return time.mktime(self.TimeTuple())
    
    def TimeTuple(self):
        return time.strptime(self.filename, 'dvgrab-%Y.%m.%d_%H-%M-%S.avi')

    def TimestampFilename(self):
        return time.strftime('date-%Y.%m.%d_%H-%M-%S.dv', self.TimeTuple())

    def SizeInBytes(self):
        return os.stat(self.filename).st_size

    def CreateTimestampSegment(self):
        timestamp = time.strftime('%d %B %Y\n%H:%M', self.TimeTuple())
        # Remove the zero-fill on single digit days of the month.
        if timestamp[0] == '0':
            timestamp = timestamp[1:]
            assert timestamp[0] != '0', self
        tmpFilename = '/tmp/DvToKino.png'
        call("convert -size 720x576 xc:black -pointsize 60 -fill white -annotate +70+150 '%s' %s" % (timestamp, tmpFilename))
        call("image2raw -r %d %s > %s" % (_numFramesForDate, tmpFilename, self.TimestampFilename()))

class DvFiles(list):
    def __init__(self, grabStartTime, *args, **kwargs):
        list.__init__(self, *args, **kwargs)
        self.sort(key = lambda dvFile: dvFile.filename)
        self.grabStartTime = grabStartTime
        self.analyzed = False
        self.timestampsCreated = False
        
    def __repr__(self):
        return 'DvFiles(%s, %s)' % (repr(self.__dict__), list.__repr__(self))
    
    def ValidDvFiles(self):
        return [dvFile for dvFile in self if dvFile.isValid]

    @staticmethod
    def Load(grabStartTime):
        return DvFiles(grabStartTime, (DvFile(filename) for filename in os.listdir(os.getcwd()) if filename.startswith('dvgrab-')))

    def Analyze(self):
        # Ignore segments timestamped using system time rather than
        # recording time, or containing too few frames.
        for dvFile in self:
            dvFile.isValid = (dvFile.Time() < self.grabStartTime and dvFile.NumFrames() > 40)
        
        # Mark the segments that should be preceded by a timestamp
        # segment. Avoid doing this for a segment that occurs less
        # than an hour after the previous one.
        dvFiles = self.ValidDvFiles()
        dvFiles[0].needsTimestamp = True
        lastDvFile = dvFiles[0]
        for dvFile in dvFiles[1:]:
            delta = dvFile.Time() - lastDvFile.Time()
            assert delta > 0, (dvFile, lastDvFile)
            dvFile.needsTimestamp = (delta > 3600)
            lastDvFile = dvFile
        self.analyzed = True

    def CreateTimestampSegments(self):
        assert self.analyzed
        dvFiles = [dvFile for dvFile in self if dvFile.isValid and dvFile.needsTimestamp]
        for dvFile in dvFiles:
            dvFile.CreateTimestampSegment()
        self.timestampsCreated = True

    def CreateSmil(self):
        with file('Kino.smil', 'w') as f:
            f.write('<?xml version="1.0"?>\n')
            f.write('<smil xmlns:smil2="http://www.w3.org/2001/SMIL20/Language">\n')
            for dvFile in self.ValidDvFiles():
                if dvFile.needsTimestamp:
                    f.write(DvFiles.__SmilSeq(dvFile.TimestampFilename(), _numFramesForDate))
                f.write(DvFiles.__SmilSeq(dvFile.filename, dvFile.NumFrames()))
            f.write('</smil>')

    def CreateDvd(self):
        assert self.timestampsCreated
        if driveSpaceInGb(os.getcwd()) < 10:
            raise Error('Insufficient drive space available to create a DVD.')
        dvdPath = 'dvd'
        if os.path.exists(dvdPath):
            raise Error("DVD path '%s' already exists." % os.path.abspath(dvdPath))
        # Convert to DVD-compatible mpeg.
        for dvFile in self.ValidDvFiles():
            if dvFile.needsTimestamp:
                call(DvFiles.__TovidCmdline(dvFile.TimestampFilename()))
            call(DvFiles.__TovidCmdline(dvFile.filename))
        self.__DvdAuthor(dvdPath)

    def __DvdAuthor(self, dvdPath):
        dvdauthorFile = 'dvdauthor.xml'
        with file(dvdauthorFile, 'w') as f:
            f.write('<dvdauthor><vmgm /><titleset><titles><pgc>\n')
            for dvFile in self.ValidDvFiles():
                if dvFile.needsTimestamp:
                    f.write(DvFiles.__DvdAuthorVob(dvFile.TimestampFilename()))
                f.write(DvFiles.__DvdAuthorVob(dvFile.filename))
            f.write('</pgc></titles></titleset></dvdauthor>\n')
        call('dvdauthor -o %s -x %s' % (dvdPath, dvdauthorFile))
        
    @staticmethod
    def __DvdAuthorVob(filename):
        return '<vob file="%s.mpg" />\n' % filename

    @staticmethod
    def __TovidCmdline(filename):
        return 'tovid -noask -pal -dvd -quality 7 -interlaced_bf -in %s -out %s' % (filename, filename)
    
    @staticmethod
    def __SmilSeq(filename, numFrames):
        return '<seq><video src="%s" clipBegin="0" clipEnd="%d"/></seq>\n' % (os.path.abspath(filename), numFrames - 1)

class State(object):
    filename = 'state'
    
    def __init__(self):
        self.grabDone = False
        self.grabStartTime = None

    def __repr__(self):
        return repr(self.__dict__)

    def __Validate(self):
        if self.grabDone and self.grabStartTime is None:
            raise Error('Invalid state.', self)
        
    def Save(self):
        self.__Validate()
        with file(State.filename, 'w') as f:
            pickle.dump(self, f)

    @staticmethod
    def Load():
        if not os.path.exists(State.filename):
            return State()
        with file(State.filename) as f:
            state = pickle.load(f)
            state.__Validate()
            return state

class Program(object):
    def __init__(self):
        self.state = State.Load()
        self.dvFiles = None

    def Clean(self):
        call('rm -rf date-*.dv date-*.dv.mpg dvd dvdauthor.xml dvgrab-*.avi.mpg Kino.smil')
        
    def Grab(self):
        if self.state.grabDone:
            raise Error('A dvgrab has already been performed in this directory.')
        if driveSpaceInGb(os.getcwd()) < 22:
            raise Error('Insufficient drive space available to perform a dvgrab.')
        # Don't save state until the dvgrab completes.
        self.state.grabStartTime = time.time()
        try:
            call("dvgrab -format dv2 -opendml -autosplit -timestamp -size 0")
        except KeyboardInterrupt:
            # The dvgrab needs to be cancelled manually by the user.
            self.state.grabDone = True
            self.state.Save()

    def AnalyzeGrab(self):
        if not self.state.grabDone:
            raise Error('A dvgrab has not been performed in this directory.')        
        self.dvFiles = DvFiles.Load(self.state.grabStartTime)
        self.dvFiles.Analyze()
        # Report.
        for dvFile in self.dvFiles:
            log(dvFile.ReportString())

    def CreateTimestampSegments(self):
        self.dvFiles.CreateTimestampSegments()

    def CreateSmil(self):
        self.dvFiles.CreateSmil()

    def CreateDvd(self):
        self.dvFiles.CreateDvd()
        
def main(argv):
    directory = argv[1]
    actions = argv[2:]

    os.chdir(directory)
    
    program = Program()
    if 'clean' in actions:
        program.Clean()
    if 'grab' in actions:
        program.Grab()
    if 'analyze' in actions or 'smil' in actions or 'dvd' in actions:
        program.AnalyzeGrab()
    if 'smil' in actions or 'dvd' in actions:
        program.CreateTimestampSegments()
        program.CreateSmil()
    if 'dvd' in actions:
        program.CreateDvd()

if __name__ == '__main__':
    main(sys.argv)
